#本脚本实现了遍历一个文件夹下所有图片并展示，可以采用键盘进行快速遍历
#a ⬅ 打开上一张图片 d ➡ 打开下一张图片 e 打开当前图片之后的第十张图片
#q打开当前图片之前的第十张图片 0将当前图片剪切到目标文件夹路径  1将当前图片剪切到目标文件夹路径
#按t切换 鼠标左键标记当前图片中的像素点位置  拖拉鼠标左键画矩形框(显示左上角、右上角坐标) 鼠标右键清除当前图片中所有框
#当前图片对应的标注文件显示(待完成)
#按r切换 涂黑任意区域，并保存图片文件至目标文件


from pathlib import Path
import argparse
import cv2
import shutil
import os
import numpy as np

def parse_arguments():
    parser = argparse.ArgumentParser(description="moveImg")
    parser.add_argument('--source_img_dir', type=str, default=r"D:\data\研究生\课题资料\服务区\数据集\output_folder")
    parser.add_argument('--target_img_dir', type=str, default=r"D:\tempDownload\dataset\elevator\nolabel_rope")
    parser.add_argument('--move_img_dir_0', type=str, default=r"D:\data\研究生\课题资料\服务区\数据集\imgs\security") #数字0对应的文件夹
    parser.add_argument('--move_img_dir_1', type=str, default=r"D:\data\研究生\课题资料\服务区\数据集\imgs\visitor") #数字1对应的文件夹
    return parser.parse_args()

def read_image_from_chinese_path(chinese_path):
    # 使用numpy从中文路径读取数据
    file_stream = open(chinese_path, 'rb')
    file_bytes = np.asarray(bytearray(file_stream.read()), dtype=np.uint8)
    file_stream.close()

    # 使用cv2.imdecode读取图片数据
    img = cv2.imdecode(file_bytes, cv2.IMREAD_COLOR)
    return img
def Draw_And_Get_Pixel_Value(event, x, y, flags, param=None):
    #原图矩阵 原图拷贝矩阵(浅拷贝) 原图拷贝中间矩阵(浅拷贝) 原图路径 左键是否按下 是否输出图片像素值
    if event == cv2.EVENT_LBUTTONDOWN:
        param['init_point'] = (x,y)
        xy = f"{param['init_point']}" if param['outputType'] else f"{param['init_point']},{str(param['image'][y][x])}"
        cv2.circle(param['image_rem'], (x, y), 5, (231, 106, 141), thickness=-1)
        cv2.putText(param['image_rem'], xy, (x, y), cv2.FONT_HERSHEY_PLAIN,
                    1.5, (25, 146, 246), thickness=2)
        param['image_mid'] = param['image_rem'].copy()
        param['lbutton_state'] = 1
    # 按下左键并滑动，不松开就持续触发
    elif param['lbutton_state'] == 1 and flags == cv2.EVENT_FLAG_LBUTTON :
        cur_point = (x,y)
        xy = f"{param['init_point']}" if param['outputType'] else f"{param['init_point']},{str(param['image'][y][x])}"
        if not (event == cv2.EVENT_LBUTTONUP): # 左键未松开，一直被清除
            #重新显示矩阵
            param['image_rem'] = param['image_mid'].copy() # 不能直接赋值操作，会直接认为是同一地址的数据
        if param['lbutton_state'] == 1:
            cv2.rectangle(param['image_rem'], param['init_point'], cur_point, (231, 106, 141), 2)
            cv2.putText(param['image_rem'],xy, (x, y), cv2.FONT_HERSHEY_PLAIN,
                    1.5, (25, 146, 246), thickness=2)
        if event == cv2.EVENT_LBUTTONUP:  # 松开左键
            param['lbutton_state'] = 0
            param['image_mid'] = param['image_rem'].copy()
    elif event == cv2.EVENT_RBUTTONDOWN:  # 右键按下执行的动作
        param['image_rem'] = param['image'].copy()

    cv2.imshow(f"{param['image_path'].name}", param['image_rem'])  # 显示的是图片矩阵

def Draw_Mask(event, x, y, flags, param=None):
    if event == cv2.EVENT_LBUTTONDOWN:
        param['lbutton_state'] = 1
        param['init_point'] = (x,y)
        param['previous_image_states'].append(param['image_rem'].copy())
    elif param['lbutton_state'] == 1 and flags == cv2.EVENT_FLAG_LBUTTON:
        cv2.line(param['image_rem'], param['init_point'], (x, y), (0, 0, 0), param["draw_size"])
        param['init_point'] = (x,y)
    elif event == cv2.EVENT_LBUTTONUP:
        param['lbutton_state'] = 0
        cv2.line(param['image_rem'], param['init_point'], (x, y), (0, 0, 0), param["draw_size"])
        cv2.imwrite(f"{param['target_img_dir']}/{param['image_path'].stem}_mask.jpg", param['image_rem'])
        print(f"saveMaskSuccess!!! 放在{param['target_img_dir']}{param['image_path'].stem}_mask.jpg")
    elif event == cv2.EVENT_RBUTTONDOWN:
        if param['previous_image_states']:
            param['image_rem'] = param['previous_image_states'].pop()
        cv2.imwrite(f"{param['target_img_dir']}/{param['image_path'].stem}_mask.jpg", param['image_rem'])
        print(f"saveMaskSuccess!!! 放在{param['target_img_dir']}{param['image_path'].stem}_mask.jpg")

    cv2.imshow(f"{param['image_path'].name}", param['image_rem'])  # 显示的是图片矩阵

def Get_Draw_Size(x, params):
    params['draw_size'] = x



def Handle_key_Press(key, index, total_images, image, image_path, target_img_dir):
    #每次读入一个有效的键盘都会进行画面的刷新，后续或可以进行优化，现在保存mask的逻辑是自动保存
    params = {
        'draw-size':20, #画笔尺寸
        'outputType':0, #矩形框模式下的输出格式 1输出位置 0输出位置加像素值
        'init_point':(0,0), #初始鼠标位置点  矩形拉框  画笔涂抹都需要
        'lbutton_state': 0, #左键是否被按下
        'previous_image_states': [],    #存储矩阵序列，可以撤回
        'image_mid': image.copy(),      #存储中间矩阵，其实可以优化掉
        'image_rem': image.copy(),      #存储备份矩阵，画图都是使用该矩阵
        'image': image,                 #原图矩阵
        'image_path': image_path,       #原图文件位置
        'target_img_dir':target_img_dir #保存文件的目标位置
    }
    if key == ord('d') or key == 2555904: #右方向键对应的ASCALL码值
        index += 1
        if index >= total_images:
            index = 0  # 回到第一张图片
            print("遍历到最后一张图片")
    elif key == ord('a') or key == 2424832:  #左方向键对应的ASCALL码值
        index -= 1
        if index < 0:
            index = total_images - 1  # 跳转到最后一张图片
            print("遍历到第一张图片")
    elif key == ord('e'):
        index += 10
        if index >= total_images:
            index = 0
            print("遍历到最后一张图片")
    elif key == ord('q'):
        index -= 10
        if index < 0:
            print("遍历到第一张图片")
            index = 0
    elif key == 27:  # 如果按下ESC，退出
        index = -1
    elif key == 48: # 小键盘数字1
        index = -2
    elif key == 49: # 小键盘数字2
        index = -3
    elif key == 3014656:  # delete
        index = -4
    elif key == ord('t'):
        cv2.destroyWindow(f'{image_path.name}')
        Creat_Window(image_path, False)
        cv2.setMouseCallback(f'{image_path.name}', Draw_And_Get_Pixel_Value, params)
        print("switch Draw_And_Get_Pixel_Value")
    elif key == ord('r'):
        cv2.createTrackbar('draw_size', f'{image_path.name}', 10, 255, lambda x:Get_Draw_Size(x, params))
        cv2.setMouseCallback(f'{image_path.name}', Draw_Mask, params)
        print("switch Draw_Mask")
    return index

def Creat_Window(image_path, resize=True):
    cv2.namedWindow(f'{image_path.name}', cv2.WINDOW_NORMAL)
    cv2.moveWindow(f'{image_path.name}', 600, 250)
    image = read_image_from_chinese_path(str(image_path))


    if resize:
        cv2.resizeWindow(f'{image_path.name}', 1500, 980)
    else:
        height, width = image.shape[:2]
        cv2.resizeWindow(f'{image_path.name}', 64, 128)

    cv2.imshow(f'{image_path.name}', image)

    return image, image_path

def main():
    #读取参数
    args = parse_arguments()
    #文件后缀列表
    extensions = ['*.jpg', '*.png']
    source_img_dir = [file for ext in extensions for file in Path(args.source_img_dir).glob(ext)]
    target_img_dir = args.target_img_dir
    move_img_dir = []
    move_img_dir.append(args.move_img_dir_0)
    move_img_dir.append(args.move_img_dir_1)
    #获得该目录下的所有jpg文件
    files = [f for f in source_img_dir]
    files_length = len(files)
    index = 0

    while True:
        #读取图片并显示
        image, image_path = Creat_Window(files[index], False)
        # 处理输入键盘逻辑
        key=cv2.waitKeyEx(0)
        new_index = Handle_key_Press(key, index, files_length, image, image_path, target_img_dir)
        #如果当前图片下标发生变化，则摧毁当前窗口
        if index != new_index:
            cv2.destroyAllWindows()
        if new_index == -1:
            break
        elif new_index == -2: #执行文件移动操作
            shutil.move(str(files[index]), move_img_dir[0])
            files.pop(index)
            print(f"moveSuccess!!! from {str(files[index])} To {str(move_img_dir[0])}")
            new_index = index if (index) < len(files) else 0
        elif new_index == -3: #执行文件移动操作
            shutil.move(str(files[index]), move_img_dir[1])
            files.pop(index)
            print(f"moveSuccess!!! from {str(files[index])} To {str(move_img_dir[1])}")
            new_index = index if (index) < len(files) else 0
        elif new_index == -4: #执行文件删除操作
            os.remove(str(files[index]))
            files.pop(index)
            print(f"deleteSuccess!!! from {str(files[index])}")
            new_index = index if (index) < len(files) else 0
        #更新下标
        index = new_index


if __name__ == "__main__":
    main()
